# coding=utf-8
### 消费者
import json
import os
import redis
from utils.oss_file import download_file_to_temp_folder
from omegaconf import OmegaConf
from dataset.data_factory import classify_provider, provider, unsupervised_provider
from train_classify import TrainVal as ClassifyTrainVal
from train_segment import TrainVal as SegmentTrainVal
from train_unsupervised import TrainVal as UnsupervisedTrainVal
from train_object_detect import TrainVal as ObjectDetectTrainVal
import socket
import torch
from utils.mysql_crud import create_connection, update_data
from utils.errors import LossFucOrganizationError

"""                       
rabbitmq消息体：
body = {
    'type': 'classify',
    'train_task_id': 1,
    'yaml_path': 'http://msjava.oss-cn-qingdao.aliyuncs.com/classify_config.yaml'
}
"""


def call_for_train_classify(config):
    train_loader, valid_loader = classify_provider(
        config.dataset.type,
        config.dataset.data_folder,
        config.batch_size,
        config.num_workers,
        config.dataset.height,
        config.dataset.width,
        config.dataset.mean,
        config.dataset.std,
        config.augmentation_flag
    )
    train_val = ClassifyTrainVal(config, device)
    train_val.train(train_loader, valid_loader)


"""
rabbitmq消息体：
body = {
    'type': 'segment',
    'train_task_id': 1,
    'yaml_path': 'http://msjava.oss-cn-qingdao.aliyuncs.com/classify_config.yaml'
}
"""


def call_for_train_segment(config):
    train_loader, valid_loader = provider(
        config.dataset.type,
        config.dataset.data_folder,
        config.batch_size,
        config.num_workers,
        config.dataset.height,
        config.dataset.width,
        config.dataset.mean,
        config.dataset.std,
        config.augmentation_flag,
        config.mask_only
    )
    train_val = SegmentTrainVal(config, device)
    train_val.train(train_loader, valid_loader)


"""
rabbitmq消息体：
body = {
    'type': 'unsupervised',
    'train_task_id': 3,
    'yaml_path': 'http://msjava.oss-cn-qingdao.aliyuncs.com/unsupervised_config.yaml'
}
"""


def call_for_train_unsupervised(config):
    train_loader, valid_loader = unsupervised_provider(
        config.dataset.type,
        config.dataset.data_folder,
        config.batch_size,
        config.num_workers,
        config.dataset.height,
        config.dataset.width,
        config.dataset.mean,
        config.dataset.std,
        config.augmentation_flag
    )
    train_val = UnsupervisedTrainVal(config, device)
    train_val.train(train_loader, valid_loader)


def call_for_train_object_detect(config):
    train_val = ObjectDetectTrainVal(config, device)
    train_val.train()


# 回调函数
def do_work(body):
    body = OmegaConf.create(json.loads(body))
    try:
        yaml_path = body.yaml_path
        local_yaml_path = download_file_to_temp_folder(yaml_path, temp_path)
        config = OmegaConf.load(local_yaml_path)
        config = OmegaConf.merge(config, body)
        ###############################
        # config.epoch = 2
        # config.batch_size = 2
        ###############################
        os.remove(local_yaml_path)

        # 前置拦截设置损失函数
        if len(config.loss.names) == 1 and config.loss.names[0] == 'auto':
            if config.type == 'classify':
                config.loss.names = ['classify_loss']
                config.loss.parameters = [{}]
            elif config.type == 'segment':
                config.loss.names = ['bce_with_logits_loss']
                config.loss.parameters = [{}]
            elif config.type == 'unsupervised':
                config.loss.names = ['mse_loss']
                config.loss.parameters = [{}]
            elif config.type == 'object-detect':
                # todo 截止目前，本项目训练yolo模型只能选用ultralytics框架默认的，扩展性不好
                pass
            config.loss.weights = [1]

        else:
            for l_n in config.loss.names:
                if l_n == 'auto':
                    raise LossFucOrganizationError()

        if config.type == 'classify':
            call_for_train_classify(config)
        elif config.type == 'segment':
            call_for_train_segment(config)
        elif config.type == 'unsupervised':
            call_for_train_unsupervised(config)
        elif config.type == 'object-detect':
            call_for_train_object_detect(config)
    except Exception as e:
        print(e)
        print("Save the exception information to the database.")
        update_query = "UPDATE train_task SET exception = %s WHERE id = %s"
        update_data(create_connection(), update_query, (str(e), body.train_task_id))


if __name__ == '__main__':
    from utils.redis_crud import get_redis_client

    temp_path = './temp'
    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    hostname = socket.gethostname()
    ip_address = socket.gethostbyname(hostname)
    # queue_name = f'ip:{ip_address}'
    queue_name = 'ip:192.168.1.180'

    client = get_redis_client()
    while True:
        message = client.blpop(queue_name, timeout=0)
        if message:
            _, msg = message
            do_work(msg)
